package com.mm.cloth.rmbg;

import android.content.Context;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;



import java.io.IOException;
import java.io.InputStream;

/**
 * 基于ONNX Runtime的MODNet抠图处理器
 * 使用modnet.onnx模型实现AI抠图功能
 */
public class MODNetProcessor {
    private static final String TAG = "MODNetProcessor";
    
    private Context context;
    private ONNXInferenceHelper onnxHelper;
    private boolean isInitialized = false;
    
    // 模型文件路径
    private static final String MODEL_PATH = "modnet.onnx";
//    private static final String MODEL_PATH = "u2net_cloth_seg.onnx";

    // 处理时间统计
    private float inferenceTime = 0;
    private float postprocessTime = 0;
    
    public MODNetProcessor(Context context) {
        this.context = context;
        initializeProcessor();
    }
    
    /**
     * 初始化处理器
     */
    private void initializeProcessor() {
        try {
            Log.d(TAG, "初始化基于ONNX Runtime的MODNet处理器");
            
            // 检查内存状态
            checkMemoryStatus();
            
            // 初始化ONNX推理助手
            onnxHelper = new ONNXInferenceHelper();
            
            // 加载MODNet模型
            loadMODNetModel();
            
            isInitialized = true;
            Log.d(TAG, "MODNet处理器初始化成功");
            
        } catch (Exception e) {
            Log.e(TAG, "MODNet处理器初始化失败", e);
            isInitialized = false;
        }
    }
    
    /**
     * 检查内存状态
     */
    private void checkMemoryStatus() {
        Runtime runtime = Runtime.getRuntime();
        long maxMemory = runtime.maxMemory();
        long totalMemory = runtime.totalMemory();
        long freeMemory = runtime.freeMemory();
        long usedMemory = totalMemory - freeMemory;
        long availableMemory = maxMemory - usedMemory;
        
        Log.d(TAG, "=== 内存状态检查 ===");
        Log.d(TAG, "最大内存: " + (maxMemory / (1024*1024)) + " MB");
        Log.d(TAG, "已用内存: " + (usedMemory / (1024*1024)) + " MB");
        Log.d(TAG, "可用内存: " + (availableMemory / (1024*1024)) + " MB");
        
        if (availableMemory < 200 * 1024 * 1024) {
            Log.w(TAG, "⚠️ 可用内存不足200MB，MODNet模型加载可能有风险");
        }
    }
    
    /**
     * 加载MODNet模型
     */
    private void loadMODNetModel() {
        try {
            Log.d(TAG, "加载MODNet ONNX模型");
            
            AssetManager assetManager = context.getAssets();
            
            // 检查模型文件是否存在
            try {
                InputStream modelStream = assetManager.open(MODEL_PATH);
                long fileSize = modelStream.available();
                modelStream.close();
                
                Log.d(TAG, "MODNet模型文件大小: " + (fileSize / (1024*1024)) + " MB");
                
                // 初始化ONNX Runtime模型
                onnxHelper.initModel(assetManager, MODEL_PATH);
                
                Log.d(TAG, "MODNet ONNX模型加载成功");
                
            } catch (IOException e) {
                Log.e(TAG, "MODNet模型文件不存在: " + MODEL_PATH);
                Log.e(TAG, "请确保将modnet.onnx文件放置在assets目录中");
                throw e;
            }
            
        } catch (Exception e) {
            Log.e(TAG, "MODNet模型加载失败", e);
            throw new RuntimeException("MODNet模型加载失败", e);
        }
    }
    
    /**
     * 执行MODNet背景移除
     * @param inputBitmap 输入图像
     * @return 处理后的图像
     */
    public Bitmap removeBackground(Bitmap inputBitmap) {
        return removeBackground(inputBitmap, null);
    }
    
    /**
     * 执行MODNet背景移除（带背景图像）
     * @param inputBitmap 输入图像
     * @param backgroundBitmap 背景图像（可为null，将生成透明背景）
     * @return 处理后的图像
     */
    public Bitmap removeBackground(Bitmap inputBitmap, Bitmap backgroundBitmap) {
        if (!isInitialized) {
            Log.e(TAG, "MODNet处理器未初始化");
            return null;
        }
        
        if (inputBitmap == null) {
            Log.e(TAG, "输入图像不能为空");
            return null;
        }
        
        try {
            Log.d(TAG, "开始MODNet背景移除处理");
            long totalStartTime = System.currentTimeMillis();

            // 检查内存状态
            System.gc(); // 建议垃圾回收

            // 准备处理

            // 1. 执行ONNX推理
            long inferenceStart = System.currentTimeMillis();
            float[] maskData = onnxHelper.predict(inputBitmap);
            long inferenceEnd = System.currentTimeMillis();
            inferenceTime = inferenceEnd - inferenceStart;

            if (maskData == null) {
                Log.e(TAG, "MODNet推理失败");
                return null;
            }

            // 2. 后处理：生成最终图像
            long postprocessStart = System.currentTimeMillis();

            // 如果没有提供背景图像，创建透明背景
            Bitmap bgImage = backgroundBitmap;
            if (bgImage == null) {
                bgImage = createTransparentBackground(inputBitmap.getWidth(), inputBitmap.getHeight());
            }

            // 合成最终图像
            Bitmap result = createMattingResult(inputBitmap, maskData, bgImage);

            // 清理临时背景图像
            if (backgroundBitmap == null && bgImage != null) {
                bgImage.recycle();
            }

            long postprocessEnd = System.currentTimeMillis();
            postprocessTime = postprocessEnd - postprocessStart;

            long totalEndTime = System.currentTimeMillis();
            Log.d(TAG, "MODNet背景移除完成，总耗时: " + (totalEndTime - totalStartTime) + "ms");
            Log.d(TAG, "时间分解 - 推理: " + inferenceTime + "ms, 后处理: " + postprocessTime + "ms");

            return result;

        } catch (OutOfMemoryError e) {
            Log.e(TAG, "MODNet背景移除内存不足", e);
            System.gc(); // 强制垃圾回收
            return null;
        } catch (Exception e) {
            Log.e(TAG, "MODNet背景移除失败", e);
            return null;
        }
    }
    
    /**
     * 创建透明背景
     */
    private Bitmap createTransparentBackground(int width, int height) {
        Bitmap background = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        background.eraseColor(Color.TRANSPARENT);
        return background;
    }
    
    /**
     * 创建抠图结果
     * @param originalImage 原始图像
     * @param maskData 遮罩数据
     * @param background 背景图像
     * @return 合成后的图像
     */
    private Bitmap createMattingResult(Bitmap originalImage, float[] maskData, Bitmap background) {
        Log.d(TAG, "开始创建抠图结果");
        
        int width = originalImage.getWidth();
        int height = originalImage.getHeight();
        
        // 将遮罩数据转换为bitmap
        int maskSize = (int) Math.sqrt(maskData.length);
        Bitmap maskBitmap = createMaskBitmap(maskData, maskSize);
        
        // 调整遮罩尺寸到原始图像尺寸
        Bitmap resizedMask = Bitmap.createScaledBitmap(maskBitmap, width, height, true);
        
        // 调整背景尺寸
        Bitmap resizedBackground = Bitmap.createScaledBitmap(background, width, height, true);
        
        // 创建结果图像
        Bitmap result = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        
        // 获取像素数据
        int[] originalPixels = new int[width * height];
        int[] backgroundPixels = new int[width * height];
        int[] maskPixels = new int[width * height];
        int[] resultPixels = new int[width * height];
        
        originalImage.getPixels(originalPixels, 0, width, 0, 0, width, height);
        resizedBackground.getPixels(backgroundPixels, 0, width, 0, 0, width, height);
        resizedMask.getPixels(maskPixels, 0, width, 0, 0, width, height);
        
        // 逐像素合成
        int transparentCount = 0;
        for (int i = 0; i < originalPixels.length; i++) {
            // 提取原始像素
            int originalPixel = originalPixels[i];
            int originalR = Color.red(originalPixel);
            int originalG = Color.green(originalPixel);
            int originalB = Color.blue(originalPixel);
            
            // 提取背景像素
            int backgroundPixel = backgroundPixels[i];
            int backgroundR = Color.red(backgroundPixel);
            int backgroundG = Color.green(backgroundPixel);
            int backgroundB = Color.blue(backgroundPixel);
            
            // 提取遮罩alpha值
            int maskPixel = maskPixels[i];
            float alpha = Color.red(maskPixel) / 255.0f;
            
            // 合成像素
            int resultR, resultG, resultB, resultA;
            
            if (alpha < 0.1f) {
                // 背景区域：使用背景像素或透明
                if (Color.alpha(backgroundPixel) == 0) {
                    resultPixels[i] = Color.TRANSPARENT;
                    transparentCount++;
                } else {
                    resultPixels[i] = backgroundPixel;
                }
            } else {
                // 前景区域：按alpha混合
                resultR = (int) (originalR * alpha + backgroundR * (1 - alpha));
                resultG = (int) (originalG * alpha + backgroundG * (1 - alpha));
                resultB = (int) (originalB * alpha + backgroundB * (1 - alpha));
                resultA = Math.max((int) (255 * alpha), Color.alpha(backgroundPixel));
                
                resultPixels[i] = Color.argb(resultA, resultR, resultG, resultB);
            }
        }
        
        result.setPixels(resultPixels, 0, width, 0, 0, width, height);
        
        // 清理临时bitmap
        try {
            if (maskBitmap != null && !maskBitmap.isRecycled()) {
                maskBitmap.recycle();
            }
            if (resizedMask != null && !resizedMask.isRecycled()) {
                resizedMask.recycle();
            }
            if (resizedBackground != background && resizedBackground != null && !resizedBackground.isRecycled()) {
                resizedBackground.recycle();
            }
        } catch (Exception e) {
            Log.w(TAG, "清理临时bitmap时出错", e);
        }

        Log.d(TAG, "抠图结果创建完成，透明像素: " + transparentCount + "/" + resultPixels.length +
              " (" + String.format("%.1f", transparentCount * 100.0f / resultPixels.length) + "%)");

        return result;
    }
    
    /**
     * 创建遮罩bitmap
     */
    private Bitmap createMaskBitmap(float[] maskData, int size) {
        Bitmap maskBitmap = Bitmap.createBitmap(size, size, Bitmap.Config.ARGB_8888);
        int[] maskPixels = new int[size * size];
        
        for (int i = 0; i < maskData.length; i++) {
            float value = Math.max(0.0f, Math.min(1.0f, maskData[i]));
            int grayValue = (int) (value * 255);
            maskPixels[i] = Color.argb(255, grayValue, grayValue, grayValue);
        }
        
        maskBitmap.setPixels(maskPixels, 0, size, 0, 0, size, size);
        return maskBitmap;
    }
    
    /**
     * 获取推理时间
     */
    public float inferenceTime() {
        return inferenceTime;
    }

    /**
     * 获取后处理时间
     */
    public float postprocessTime() {
        return postprocessTime;
    }
    
    /**
     * 释放资源
     */
    public void release() {
        try {
            if (onnxHelper != null) {
                onnxHelper.release();
                onnxHelper = null;
            }

            // 强制垃圾回收
            System.gc();

            isInitialized = false;
            Log.d(TAG, "MODNet处理器资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放MODNet资源失败", e);
        }
    }
    
    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return isInitialized;
    }
}
